In [2]:
import networkx as nx
Define as currently and check working
In [3]:
def test__rows_to_columns(test_func):
G = nx.Graph()
G.add_nodes_from([(0, {"attr_1": "a", "attr_2": 10}),
(1, {"attr_1": "b"}),
(2, {"attr_1": "c", "attr_2": 30})])
G.add_edges_from([(0, 1, {"attr_1": "A"}),
(0, 2, {"attr_1": "B", "attr_2": 10})])
node_attr_keys = ["attr_1", "attr_2"]
node_dict = test_func(G.nodes(data=True), node_attr_keys)
assert node_dict["attr_1"] == [attr['attr_1'] for _, attr in G.nodes(data=True)], "returned: %s" % node_dict['attr_1']
assert node_dict["attr_2"] == [attr['attr_2'] if key != 1 else None for key, attr in G.nodes(data=True)]
edge_attr_keys = ["attr_1", "attr_2"]
edge_dict = test_func(G.edges(data=True), edge_attr_keys)
assert edge_dict["attr_1"] == ["A", "B"]
assert edge_dict["attr_2"] == [None, 10]
In [4]:
def _rows_to_columns_orig(source, attr_keys):
attr_dict = {}
for attr_key in attr_keys:
attr_dict[attr_key] = [attr[attr_key] if attr_key in attr.keys() else None for *_, attr in source]
return attr_dict
In [ ]:
test__rows_to_columns(_rows_to_columns_orig)
Redefine in py2/3 syntax
In [6]:
def _rows_to_columns_23(nx_data, attr_keys):
# nx_data may be coming from edges or nodes
# So extract the data depending on the type
if isinstance(nx_data, nx.classes.reportviews.EdgeDataView):
source = [i[2] for i in nx_data]
elif isinstance(nx_data, nx.classes.reportviews.NodeDataView):
source = [data for node, data in nx_data]
attr_dict = {}
for attr_key in attr_keys:
compiled_data = []
for data in source:
compiled_data.append(data.get(attr_key))
attr_dict[attr_key] = compiled_data
return attr_dict
In [7]:
test__rows_to_columns(_rows_to_columns_23)
A quick look at the inputs to understand what's happening here
In [8]:
G = nx.Graph()
G.add_nodes_from([(0, {"attr_1": "a", "attr_2": 10}),
(1, {"attr_1": "b"}),
(2, {"attr_1": "c", "attr_2": 30})])
G.add_edges_from([(0, 1, {"attr_1": "A"}),
(0, 2, {"attr_1": "B", "attr_2": 10})])
In [9]:
G.edges(data=True)
Out[9]:
In [10]:
G.nodes(data=True)
Out[10]:
In [ ]: